import asyncio
import logging
from typing import Tuple

# 配置日志格式
logging.basicConfig(
    format="[%(asctime)s] %(levelname)s: %(message)s",
    level=logging.INFO,
    datefmt="%Y-%m-%d %H:%M:%S"
)
logger = logging.getLogger(__name__)

class AsyncUDPServer:
    def __init__(self, host: str = "0.0.0.0", port: int = 8888):
        """
        异步UDP服务器
        :param host: 监听地址
        :param port: 监听端口
        """
        self.host = host
        self.port = port
        self.transport = None
        self.protocol = None
        
    async def start(self):
        """启动UDP服务器"""
        loop = asyncio.get_running_loop()
        
        # 创建UDP端点
        self.transport, self.protocol = await loop.create_datagram_endpoint(
            lambda: UDPProtocol(self.handle_datagram),
            local_addr=(self.host, self.port)
        )
        
        logger.info(f"UDP Server started on {self.host}:{self.port}")

    async def stop(self):
        """停止服务器"""
        if self.transport:
            self.transport.close()
            logger.info("UDP Server stopped")

    async def handle_datagram(self, data: bytes, addr: Tuple[str, int]):
        """
        处理接收到的UDP数据包（需被子类重写）
        :param data: 接收到的原始字节数据
        :param addr: 客户端地址 (ip, port)
        """
        # 默认实现：原样返回数据（echo服务）
        logger.debug(f"Received {len(data)} bytes from {addr[0]}:{addr[1]}")
        self.send_response(data, addr)

    def send_response(self, data: bytes, addr: Tuple[str, int]):
        """
        发送响应到指定客户端
        :param data: 要发送的字节数据
        :param addr: 目标地址 (ip, port)
        """
        if self.transport and not self.transport.is_closing():
            self.transport.sendto(data, addr)
            logger.debug(f"Sent {len(data)} bytes to {addr[0]}:{addr[1]}")

class UDPProtocol(asyncio.DatagramProtocol):
    """自定义UDP协议处理器"""
    def __init__(self, data_handler):
        super().__init__()
        self.data_handler = data_handler
        self.transport = None

    def connection_made(self, transport):
        """连接建立时调用"""
        self.transport = transport

    def datagram_received(self, data: bytes, addr: Tuple[str, int]):
        """收到数据包时调用"""
        # 将数据包处理委托给主服务器
        asyncio.create_task(self.data_handler(data, addr))

    def error_received(self, exc):
        """接收错误时调用"""
        logger.error(f"UDP error received: {exc}")

    def connection_lost(self, exc):
        """连接关闭时调用"""
        if exc:
            logger.error(f"UDP connection lost: {exc}")
        else:
            logger.info("UDP connection closed normally")

# 示例使用：带业务逻辑的自定义服务器
class CustomUDPServer(AsyncUDPServer):
    async def handle_datagram(self, data: bytes, addr: Tuple[str, int]):
        """自定义业务逻辑处理"""
        try:
            # 示例：将数据转为大写并返回
            text = data.decode("utf-8").strip()
            logger.info(f"Received from {addr}: '{text}'")
            
            response = f"ECHO: {text.upper()}\n".encode("utf-8")
            self.send_response(response, addr)
            
        except UnicodeDecodeError:
            error_msg = b"Error: Invalid UTF-8 data\n"
            self.send_response(error_msg, addr)

async def start(port: int = 8888):
    """主运行函数"""
    # echo "Hello UDP" | nc -u 127.0.0.1 8888 测试
    server = CustomUDPServer(port=port)
    await server.start()
    await asyncio.Event().wait()